"""


Version: 0.1
Author: lk
Date: 2022-03-08 13:39
"""
import dgl
import torch
from torch import nn

def moment_update(model:nn.Module,model_ema:nn.Module,m:float):
    for p1,p2 in zip(model.parameters(),model_ema.parameters()):
        p2.data.mul_(m).add((1-m) * p1.detach().data)


class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.linear = nn.Linear(3,4)
if __name__ == '__main__':
    model= MyModel()
    model_ema=MyModel()
    # moment_update(model,model_ema,0.99)

    # def set_bn_train(module):
    #     class_name= module.__class__.__name__
    #     print(module)
    #     print(class_name)
    #
    #     if class_name.find("BatchNorm") != -1:
    #         module.train()
    # model.apply(set_bn_train)

    # u,v = torch.tensor([0,1,2]),torch.tensor([1,0,0])
    # g = dgl.graph((u,v))
    # print(g)
    # zero_degree_node=(g.in_degrees()==0).nonzero().squeeze()
    # print(zero_degree_node)
    # # g.remove_nodes(zero_degree_node)
    #
